package com.matrix.app.flux.aichat.service.impl;

import com.matrix.app.flux.aichat.data.AiChatRequestDto;
import com.matrix.app.flux.aichat.data.AiChatResponseVo;
import com.matrix.app.flux.aichat.service.AiChatService;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;

/**
 * AI聊天服务实现类
 *
 * Copyright © 海平面工作室 版权所有
 *
 * @Author: Leo
 * @Create: 2024-01-22 10:00:00
 * @Since 1.2
 */
@Service
public class AiChatServiceImpl implements AiChatService {

    @Autowired
    private ChatClient chatClient;

    @Override
    public Mono<AiChatResponseVo> chat(AiChatRequestDto request) {
        return Mono.fromCallable(() -> {
            List<Message> messages = buildMessages(request);
            
            var response = chatClient.prompt()
                    .messages(messages)
                    .call()
                    .chatResponse();
            
            String content = response.getResult().getOutput().getContent();
            Integer tokensUsed = Math.toIntExact(response.getMetadata().getUsage().getTotalTokens());
            
            return new AiChatResponseVo(content, request.getModelName(), tokensUsed);
        });
    }

    @Override
    public Flux<String> streamChat(AiChatRequestDto request) {
        return Flux.create(sink -> {
            try {
                List<Message> messages = buildMessages(request);
                
                chatClient.prompt()
                        .messages(messages)
                        .stream()
                        .chatResponse()
                        .doOnNext(response -> {
                            String content = response.getResult().getOutput().getContent();
                            if (content != null && !content.isEmpty()) {
                                sink.next("data: " + content + "\n\n");
                            }
                        })
                        .doOnComplete(sink::complete)
                        .doOnError(sink::error)
                        .subscribe();
            } catch (Exception e) {
                sink.error(e);
            }
        });
    }

    private List<Message> buildMessages(AiChatRequestDto request) {
        List<Message> messages = new ArrayList<>();
        
        // 添加历史消息
        if (request.getHistory() != null && !request.getHistory().isEmpty()) {
            for (AiChatRequestDto.AiChatMessage historyMsg : request.getHistory()) {
                if ("user".equals(historyMsg.getRole())) {
                    messages.add(new UserMessage(historyMsg.getContent()));
                } else if ("assistant".equals(historyMsg.getRole())) {
                    messages.add(new AssistantMessage(historyMsg.getContent()));
                }
            }
        }
        
        // 添加当前用户消息
        messages.add(new UserMessage(request.getMessage()));
        
        return messages;
    }
}